package com.pg.agent.engine.service.utils;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.pg.agent.engine.client.common.AgentConstant;
import com.pg.agent.engine.client.component.AgentComponentLLM;
import com.pg.agent.engine.client.message.InPutMessage;
import com.pg.agent.engine.client.message.OutPutMessage;
import com.pg.agent.engine.client.model.ComponentInParam;
import com.pg.agent.engine.client.model.Function;
import com.pg.agent.engine.client.model.llm.CallLog;
import com.pg.agent.engine.client.model.llm.ComponentLLM;
import com.pg.agent.engine.client.model.llm.LLMPrompt;
import com.pg.agent.engine.client.model.llm.LLMPromptType;
import com.pg.agent.engine.client.model.sop.Memory;
import com.pg.base.component.PgAssert;
import com.pg.base.exception.BizException;
import com.pg.base.util.PgHttpUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import javax.servlet.http.HttpServletResponse;
import java.util.*;

/**
 * llm 工具类
 *
 * @author 张俊敏
 * @date 2023/11/15 16:42
 */
@Slf4j
public class ComponentLLMUtil {


    public static void call(CallLog lLMCallLog, ComponentLLM llm, InPutMessage message, List<Memory> chatHistory, String systemPrompt, String lastPrompt) {

        CallLog callInfo = getResponse(lLMCallLog, llm, chatHistory, systemPrompt, lastPrompt);

        // 输出日志
        logInfo(lLMCallLog);

        message.getCallLogList().add(callInfo);

    }

    public static void getStreamCall(CallLog lLMCallLog, ComponentLLM llm, InPutMessage message, List<Memory> chatHistory,
                                     String systemPrompt, String lastPrompt, HttpServletResponse response, OutPutMessage outMessage) {

        getStreamResponse(lLMCallLog, llm, message, chatHistory, systemPrompt, lastPrompt, response, message.getMessageId(), outMessage);

        // 输出日志
        logInfo(lLMCallLog);

    }

    private static void getStreamResponse(CallLog lLMCallLog, ComponentLLM llm, InPutMessage message, List<Memory> chatHistory,
                                          String systemPrompt, String lastPrompt, HttpServletResponse response, String messageId, OutPutMessage outMessage) {

        //对话提示语：如：请尽量简明扼要地答复,控制在100个字内
        List<LLMPrompt> promptList = new ArrayList<>();
        //promptList.add(new LLMPrompt(LLMPromptType.SYSTEM, StringUtils.trimToEmpty(systemPrompt)));

        // 获取最后10条数据
        if (CollectionUtil.isNotEmpty(chatHistory) && chatHistory.size() > AgentConstant.MAX_CHAT_HISTORY) {
            chatHistory = chatHistory.subList(chatHistory.size() - AgentConstant.MAX_CHAT_HISTORY, chatHistory.size());
        }

        if (chatHistory != null) {
            // 对历史对话记录进行倒序组装成发送给大模型数据格式
            for (Memory memory : chatHistory) {
                promptList.add(new LLMPrompt(LLMPromptType.USER, StringUtils.trimToEmpty(memory.getContent())));
            }
        }

        // 在最后追加提示语，返回什么结果提示
        if (ObjectUtil.isNotEmpty(lastPrompt)) {
            LLMPrompt userPrompt = promptList.get(promptList.size() - 1);
            userPrompt.setContent(userPrompt.getContent() + lastPrompt);
        }
        String result = null;
        try {
            JSONObject param = new JSONObject();
            param.put("messages", promptList);
            List<ComponentInParam> componentInParamList = JSONObject.parseArray(llm.getInParam(), ComponentInParam.class);

            if (componentInParamList!=null) {
                for (ComponentInParam component : componentInParamList) {
                    String componentCode = component.getParamCode();
                    if ("modelName".equals(componentCode)) {
                        param.put("modelName", "gpt-4");
                    }
                    if ("temperature".equals(componentCode)) {
                        param.put("temperature", 0.75);
                    }
                }
            }
            log.info("请求流式的参数：{}", JSONObject.toJSONString(param));
            // 流式执行
            if (response != null && llm.getType() != null && llm.getType() == 1) {

                result = chatStreamCompletions(llm, message, param, response, messageId, outMessage);
            } else {

                //非流式执行
                result = chatCompletions(llm, param);

                JSONObject jsonObject = JSON.parseObject(result);

                String llmCompany = Optional.ofNullable(llm.getLlmCompany()).orElse("");
                switch (llmCompany) {
                    case "baidu":
                        result = jsonObject.getString("result");
                        break;
                    case "openAi":
                        result = jsonObject.getString("content");
                        break;
                    default:
                        result = jsonObject.getString("result");
                        break;
                }

            }

        } catch (Exception e) {
            result = "调用模型报错：" + e.getMessage();
            log.error("调用模型报错：", e);
        }
        lLMCallLog.setLLmPromptInfo(promptList, llm, result);
        message.getCallLogList().add(lLMCallLog);
    }

    private static String chatStreamCompletions(ComponentLLM llm,  InPutMessage message, JSONObject jsonObject, HttpServletResponse response, String messageId, OutPutMessage outMessage) {
        SSEListener sseListener = new SSEListener(jsonObject, message, response, messageId, outMessage, false);
        try {
            ExecuteSSEUtil.executeSSE(llm.getUrl(), sseListener, JSON.toJSONString(jsonObject));
        } catch (Exception e) {
            throw new BizException(1000,"大模型执行错误");
        }
        return sseListener.getOutput().toString();
    }


    public static void logInfo(CallLog lLMCallLog) {

        log.info("\n大模型【{}】*******{}({})*******提示词组装如下*******\n{}\n{},\n********大模型结果如下*******\n{} \n***************************结束****************************",
                lLMCallLog.getModel(), lLMCallLog.getCallType(), lLMCallLog.getCallName(), lLMCallLog.getSystemPrompt(),
                lLMCallLog.getUserPrompt(), lLMCallLog.getResponse());
    }

    /**
     * 获取LLM请求响应数据
     *
     * @param chatHistory  历史记忆
     * @param systemPrompt 系统场景对话
     * @param lastPrompt   输出提示
     * @return 模型返回结果
     */
    private static CallLog getResponse(CallLog lLMCallLog, ComponentLLM llm, List<Memory> chatHistory, String systemPrompt, String lastPrompt) {

        //对话提示语：如：请尽量简明扼要地答复,控制在100个字内
        List<LLMPrompt> promptList = new ArrayList<>();
        promptList.add(new LLMPrompt(LLMPromptType.SYSTEM, StringUtils.trimToEmpty(systemPrompt)));

        // 获取最后10条数据
        if (CollectionUtil.isNotEmpty(chatHistory) && chatHistory.size() > AgentConstant.MAX_CHAT_HISTORY) {
            chatHistory = chatHistory.subList(chatHistory.size() - AgentConstant.MAX_CHAT_HISTORY, chatHistory.size());
        }

        if (chatHistory != null) {
            // 对历史对话记录进行倒序组装成发送给大模型数据格式
            for (Memory memory : chatHistory) {
                promptList.add(new LLMPrompt(LLMPromptType.SYSTEM, StringUtils.trimToEmpty(memory.getContent())));
            }
        }

        // 在最后追加提示语，返回什么结果提示
        if (ObjectUtil.isNotEmpty(lastPrompt)) {
            LLMPrompt userPrompt = promptList.get(promptList.size() - 1);
            userPrompt.setContent(userPrompt.getContent() + lastPrompt);
        }

        String result = null;
        try {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", promptList);
            String type = lLMCallLog.getCallType();
            if ("State切换".equals(type) || "Agent切换".equals(type) || "AgentComponent".equals(type)) {
                jsonObject.put("temperature", 0);
            }
            result = chatCompletions(llm, jsonObject);

        } catch (Exception e) {
            log.error("调用模型报错：", e);
        }
        //log.info("response：{}", response);
        // 判断模型，对应返回不同的数据结构
        JSONObject jsonObject = JSON.parseObject(result);
        if (jsonObject != null && ObjectUtil.equal(jsonObject.getInteger("code"), 0)) {
            JSONObject resultObj = jsonObject.getJSONArray("data").getJSONObject(0);
            String dataType = resultObj.getString("dataType");
            if ("image".equals(dataType)) {
                lLMCallLog.setImageList(resultObj.getJSONArray("content").toJavaList(String.class));
            } else if ("text".equals(dataType)) {
                result = resultObj.getString("content");
            } else {
                result = resultObj.getString("content");
            }
            if (resultObj.getJSONObject("functionCall") != null) {
                String argument = resultObj.getJSONObject("functionCall").getString("arguments");
                lLMCallLog.setArguments(argument);
            }
        } else {
            result = "调用模型出错：" + result;
        }

        lLMCallLog.setLLmPromptInfo(promptList, llm, result);

        return lLMCallLog;
    }

    /**
     * 构建大模型对接对象和完成对话
     *
     * @param llm
     * @param jsonObject
     * @return
     */
    private static String chatCompletions(ComponentLLM llm, JSONObject jsonObject) {

        String result = null;
        Map<String, String> heard = new HashMap<>();
        heard.put("Content-Type", "application/json");
        //调用完成对话
        try {
            if (PgHttpUtil.GET.equals(llm.getRequestMethod())) {
                result = PgHttpUtil.get(llm.getUrl(), new HashMap<>(), jsonObject, 60 * 1000);
            } else {
                result = PgHttpUtil.postBody(llm.getUrl(), new HashMap<>(), JSONObject.toJSONString(jsonObject), 60 * 1000);
            }
        } catch (Exception e) {
            String errorMsg = String.format("url[%s],param[%s]", llm.getUrl(), JSON.toJSONString(jsonObject));
            PgAssert.bizError(true, 1000, "大模型接口" + errorMsg + "请求出错,原因：" + e.getMessage());
        }
        return result;
    }

    public static void callBack(CallLog lLMCallLog, ComponentLLM llm, InPutMessage message, List<Memory> chatHistory,
                                String systemPrompt, String lastPrompt, List<Function> functionList) {
        List<LLMPrompt> promptList = new ArrayList<>();
        promptList.add(new LLMPrompt(LLMPromptType.SYSTEM, StringUtils.trimToEmpty(systemPrompt)));


        // 获取最后10条数据
        if (CollectionUtil.isNotEmpty(chatHistory) && chatHistory.size() > AgentConstant.MAX_CHAT_HISTORY) {
            chatHistory = chatHistory.subList(chatHistory.size() - AgentConstant.MAX_CHAT_HISTORY, chatHistory.size());
        }
        if (chatHistory != null) {
            // 对历史对话记录进行倒序组装成发送给大模型数据格式
            for (Memory memory : chatHistory) {
                promptList.add(new LLMPrompt(LLMPromptType.USER, StringUtils.trimToEmpty(memory.getContent())));
            }
        }

        // 在最后追加提示语，返回什么结果提示
        if (ObjectUtil.isNotEmpty(lastPrompt) && promptList.size() > 0) {
            LLMPrompt userPrompt = promptList.get(promptList.size() - 1);
            userPrompt.setContent(userPrompt.getContent() + lastPrompt);
        }

        String response = null;

        try {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", promptList);
            jsonObject.put("functions", functionList);
            response = chatCompletions(llm, jsonObject);

        } catch (Exception e) {
            log.error("调用模型报错：", e);
        }

        // 判断模型，对应返回不同的数据结构
        JSONObject jsonObject = JSON.parseObject(response);
        if (jsonObject != null && jsonObject.getInteger("code") == 0) {
            JSONObject data = jsonObject.getJSONArray("data").getJSONObject(0);
            response = data.getString("content");
            lLMCallLog.setResponse(response);
            if (data.getJSONObject("functionCall") != null) {
                String argument = data.getJSONObject("functionCall").getString("arguments");
                lLMCallLog.setArguments(argument);
                if (ObjectUtil.isNotEmpty(argument)) {
                    lLMCallLog.setResponse(argument);
                }
            }
        } else {
            log.info("第三方接口请求出错,原因：{}", response);
        }

        // 追加一下参数日志
        if (functionList != null) {
            promptList.get(0).setContent(promptList.get(0).getContent() + "\n参数：" + JSON.toJSONString(functionList));
        }

        lLMCallLog.setLLmPromptInfo(promptList, llm, response);
        message.getCallLogList().add(lLMCallLog);
        // 输出日志
        logInfo(lLMCallLog);

    }

    public static CallLog callConvert(CallLog callLog, ComponentLLM llm, InPutMessage message, String userPrompt, String systemPrompt) {

        List<LLMPrompt> promptList = new ArrayList<>();
        promptList.add(new LLMPrompt(LLMPromptType.SYSTEM, StringUtils.trimToEmpty(systemPrompt)));
        promptList.add(new LLMPrompt(LLMPromptType.USER, StringUtils.trimToEmpty(userPrompt)));

        String response = null;

        try {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", promptList);
            jsonObject.put("temperature", 0.01);
            response = chatCompletions(llm, jsonObject);
        } catch (Exception e) {
            response = e.getMessage();
            log.error("调用模型报错：", e);
        }

        callLog.setLlm(llm);
        callLog.setSystemPrompt(systemPrompt);
        callLog.setUserPrompt(userPrompt);
        message.getCallLogList().add(callLog);

        // 判断模型，对应返回不同的数据结构
        JSONObject jsonObject = JSON.parseObject(response);

        if (jsonObject != null && ObjectUtil.equal(jsonObject.getInteger("code"), 0)) {
            callLog.setResponse(response);
        } else {
            callLog.setResponse("大模型调用出错" + response);
            PgAssert.bizError(true, 1000, "大模型数据结构错误：" + response);
        }

        // 输出日志
        logInfo(callLog);

        return callLog;
    }


    public static String extractResult(String response, List<String> resultList) {

        String type = "end";
        //截取出大模型放回结果
        String targetStr = get_content_between_a_b("<" + type + ">", "</" + type + ">", response);
        if (ObjectUtil.isEmpty(targetStr) && ObjectUtil.isNotEmpty(resultList)) {
            for (String result : resultList) {
                if (response.contains(result)) {
                    targetStr = result;
                    break;
                }
            }
        }
        return targetStr;
    }

    public static String get_content_between_a_b(String a, String b, String text) {
        if (ObjectUtil.isEmpty(text)) {
            return "";
        }
        int start = text.indexOf(a);
        int end = text.indexOf(b, start + a.length());
        if (start == -1 || end == -1) {
            return "";
        } else {
            return text.substring(start + a.length(), end);
        }
    }

    public static void callAgentComponent(CallLog callLog, OutPutMessage outMessage, AgentComponentLLM agentComponentLLM, InPutMessage message, String userPrompt, String systemPrompt, HttpServletResponse response) {
        List<LLMPrompt> promptList = new ArrayList<>();
        promptList.add(new LLMPrompt(LLMPromptType.SYSTEM, StringUtils.trimToEmpty(systemPrompt)));
        promptList.add(new LLMPrompt(LLMPromptType.USER, StringUtils.trimToEmpty(userPrompt)));
        ComponentLLM llm = agentComponentLLM.getComponentLLM();
        double temperature = agentComponentLLM.getTemperature();
        String result = null;
        try {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", promptList);
            jsonObject.put("temperature", temperature);
        log.info("请求流式的参数：{}", JSONObject.toJSONString(jsonObject));
        callLog.setArguments(JSONObject.toJSONString(jsonObject));
        if (response != null && llm.getType() != null && llm.getType() == 1) {
            result = chatStreamCompletions(llm, message, jsonObject, response, message.getMessageId(), outMessage);
        } else {
            result = chatCompletions(llm, jsonObject);
            JSONObject jsonObject1 = JSON.parseObject(result);
            if (jsonObject1 != null && ObjectUtil.equal(jsonObject1.getInteger("code"), 0)) {
                JSONObject resultObj = jsonObject1.getJSONArray("data").getJSONObject(0);
                result = resultObj.getString("content");
            }
        }
        } catch (Exception e) {
            result = e.getMessage();
            log.error("调用模型报错：", e);
        }

        callLog.setLLmPromptInfo(promptList,llm,result);
        message.getCallLogList().add(callLog);
    }
}
